use base64::{prelude::BASE64_STANDARD, Engine};

// Turn String into [u32]
fn string2u32(input: String) -> Vec<u32> {
    let mut result = Vec::new();
    let mut cur_blk: u32 = 0;
    let mut cur_len = 0;

    for byte in input.as_bytes() {
        cur_blk = cur_blk << 8;
        cur_blk += *byte as u32;
        cur_len += 1;
        if cur_len == 4 {
            result.push(cur_blk);
            cur_blk = 0;
            cur_len = 0;
        }
    }

    let padding = 4 - cur_len;
    for _ in 0..(4-cur_len) {
        cur_blk = cur_blk << 8;
        cur_blk += padding as u32;
    }
    result.push(cur_blk);

    result
}


pub fn encryption(input: String) -> String {
    let bit_blocks = string2u32(input);
    let mut context = Context::new();
    zuc_init(&mut context, [0u8; 16], [1u8; 16]);
    let mut encrypted = vec![];

    let keys = zuc_keygen(&mut context, bit_blocks.len() as u32);
    for (blk, key) in bit_blocks.iter().zip(keys.iter()) {
        encrypted.extend((blk ^ key).to_be_bytes());
    }

    BASE64_STANDARD.encode(encrypted)
}

const D: [u16; 16] = [
    0x44D7, 0x26BC, 0x626B, 0x135E, 0x5789, 0x35E2, 0x7135, 0x09AF,
    0x4D78, 0x2F13, 0x6BC4, 0x1AF1, 0x5E26, 0x3C4D, 0x789A, 0x47AC
];

const S0: [u8; 256] = [
    0x3E,0x72,0x5B,0x47,0xCA,0xE0,0x00,0x33,0x04,0xD1,0x54,0x98,0x09,0xB9,0x6D,0xCB,
    0x7B,0x1B,0xF9,0x32,0xAF,0x9D,0x6A,0xA5,0xB8,0x2D,0xFC,0x1D,0x08,0x53,0x03,0x90,
    0x4D,0x4E,0x84,0x99,0xE4,0xCE,0xD9,0x91,0xDD,0xB6,0x85,0x48,0x8B,0x29,0x6E,0xAC,
    0xCD,0xC1,0xF8,0x1E,0x73,0x43,0x69,0xC6,0xB5,0xBD,0xFD,0x39,0x63,0x20,0xD4,0x38,
    0x76,0x7D,0xB2,0xA7,0xCF,0xED,0x57,0xC5,0xF3,0x2C,0xBB,0x14,0x21,0x06,0x55,0x9B,
    0xE3,0xEF,0x5E,0x31,0x4F,0x7F,0x5A,0xA4,0x0D,0x82,0x51,0x49,0x5F,0xBA,0x58,0x1C,
    0x4A,0x16,0xD5,0x17,0xA8,0x92,0x24,0x1F,0x8C,0xFF,0xD8,0xAE,0x2E,0x01,0xD3,0xAD,
    0x3B,0x4B,0xDA,0x46,0xEB,0xC9,0xDE,0x9A,0x8F,0x87,0xD7,0x3A,0x80,0x6F,0x2F,0xC8,
    0xB1,0xB4,0x37,0xF7,0x0A,0x22,0x13,0x28,0x7C,0xCC,0x3C,0x89,0xC7,0xC3,0x96,0x56,
    0x07,0xBF,0x7E,0xF0,0x0B,0x2B,0x97,0x52,0x35,0x41,0x79,0x61,0xA6,0x4C,0x10,0xFE,
    0xBC,0x26,0x95,0x88,0x8A,0xB0,0xA3,0xFB,0xC0,0x18,0x94,0xF2,0xE1,0xE5,0xE9,0x5D,
    0xD0,0xDC,0x11,0x66,0x64,0x5C,0xEC,0x59,0x42,0x75,0x12,0xF5,0x74,0x9C,0xAA,0x23,
    0x0E,0x86,0xAB,0xBE,0x2A,0x02,0xE7,0x67,0xE6,0x44,0xA2,0x6C,0xC2,0x93,0x9F,0xF1,
    0xF6,0xFA,0x36,0xD2,0x50,0x68,0x9E,0x62,0x71,0x15,0x3D,0xD6,0x40,0xC4,0xE2,0x0F,
    0x8E,0x83,0x77,0x6B,0x25,0x05,0x3F,0x0C,0x30,0xEA,0x70,0xB7,0xA1,0xE8,0xA9,0x65,
    0x8D,0x27,0x1A,0xDB,0x81,0xB3,0xA0,0xF4,0x45,0x7A,0x19,0xDF,0xEE,0x78,0x34,0x60
];

const S1: [u8; 256] = [
    0x55,0xC2,0x63,0x71,0x3B,0xC8,0x47,0x86,0x9F,0x3C,0xDA,0x5B,0x29,0xAA,0xFD,0x77,
    0x8C,0xC5,0x94,0x0C,0xA6,0x1A,0x13,0x00,0xE3,0xA8,0x16,0x72,0x40,0xF9,0xF8,0x42,
    0x44,0x26,0x68,0x96,0x81,0xD9,0x45,0x3E,0x10,0x76,0xC6,0xA7,0x8B,0x39,0x43,0xE1,
    0x3A,0xB5,0x56,0x2A,0xC0,0x6D,0xB3,0x05,0x22,0x66,0xBF,0xDC,0x0B,0xFA,0x62,0x48,
    0xDD,0x20,0x11,0x06,0x36,0xC9,0xC1,0xCF,0xF6,0x27,0x52,0xBB,0x69,0xF5,0xD4,0x87,
    0x7F,0x84,0x4C,0xD2,0x9C,0x57,0xA4,0xBC,0x4F,0x9A,0xDF,0xFE,0xD6,0x8D,0x7A,0xEB,
    0x2B,0x53,0xD8,0x5C,0xA1,0x14,0x17,0xFB,0x23,0xD5,0x7D,0x30,0x67,0x73,0x08,0x09,
    0xEE,0xB7,0x70,0x3F,0x61,0xB2,0x19,0x8E,0x4E,0xE5,0x4B,0x93,0x8F,0x5D,0xDB,0xA9,
    0xAD,0xF1,0xAE,0x2E,0xCB,0x0D,0xFC,0xF4,0x2D,0x46,0x6E,0x1D,0x97,0xE8,0xD1,0xE9,
    0x4D,0x37,0xA5,0x75,0x5E,0x83,0x9E,0xAB,0x82,0x9D,0xB9,0x1C,0xE0,0xCD,0x49,0x89,
    0x01,0xB6,0xBD,0x58,0x24,0xA2,0x5F,0x38,0x78,0x99,0x15,0x90,0x50,0xB8,0x95,0xE4,
    0xD0,0x91,0xC7,0xCE,0xED,0x0F,0xB4,0x6F,0xA0,0xCC,0xF0,0x02,0x4A,0x79,0xC3,0xDE,
    0xA3,0xEF,0xEA,0x51,0xE6,0x6B,0x18,0xEC,0x1B,0x2C,0x80,0xF7,0x74,0xE7,0xFF,0x21,
    0x5A,0x6A,0x54,0x1E,0x41,0x31,0x92,0x35,0xC4,0x33,0x07,0x0A,0xBA,0x7E,0x0E,0x34,
    0x88,0xB1,0x98,0x7C,0xF3,0x3D,0x60,0x6C,0x7B,0xCA,0xD3,0x1F,0x32,0x65,0x04,0x28,
    0x64,0xBE,0x85,0x9B,0x2F,0x59,0x8A,0xD7,0xB0,0x25,0xAC,0xAF,0x12,0x03,0xE2,0xF2
];

pub struct Context {
    s: [u32; 16],
    x: [u32; 4],
    r1: u32,
    r2: u32,
}

impl Context {
    pub fn new() -> Context {
        Context {
            s: [0; 16],
            x: [0; 4],
            r1: 0,
            r2: 0,
        }
    }
}

// 初始化 ZUC
pub fn zuc_init(context: &mut Context, k: [u8; 16], iv: [u8; 16]) {
    for i in 0..16 {
        context.s[i] = concat_u31(k[i], D[i], iv[i]);
    }
    context.r1 = 0;
    context.r2 = 0;
    for _ in 0..32 {
        bit_reconstruction(context);
        let w = nonlinear_f(context);
        lfsr_with_initialization_mode(w >> 1, &mut context.s);
    }
}

// 密钥生成算法
pub fn zuc_keygen(context: &mut Context, length: u32) -> Vec<u32> {
    let mut ans = vec![];
    bit_reconstruction(context);
    nonlinear_f(context);
    lfsr_with_work_mode(&mut context.s);
    for _ in 0..length {
        bit_reconstruction(context);
        ans.push(nonlinear_f(context) ^ context.x[3]);
        lfsr_with_work_mode(&mut context.s);
    }
    ans
}

fn concat_u31(k: u8, d: u16, iv: u8) -> u32 {
    let mut res = 0;
    res |= (k as u32) << 23;
    res |= (d as u32) << 8;
    res |= iv as u32;
    res
}

fn concat_u32(a: u8, b: u8, c: u8, d: u8) -> u32 {
    let mut res = 0;
    res |= (a as u32) << 24;
    res |= (b as u32) << 16;
    res |= (c as u32) << 8;
    res |= d as u32;
    res
}

fn bit_reconstruction(context: &mut Context) {
    context.x[0] = ((context.s[15] & 0x7FFF8000) << 1) | (context.s[14] & 0xFFFF);
    context.x[1] = ((context.s[11] & 0xFFFF) << 16) | (context.s[9] >> 15);
    context.x[2] = ((context.s[7] & 0xFFFF) << 16) | (context.s[5] >> 15);
    context.x[3] = ((context.s[2] & 0xFFFF) << 16) | (context.s[0] >> 15);
}

fn nonlinear_f(context: &mut Context) -> u32 {
    let w = (context.x[0] ^ context.r1).wrapping_add(context.r2);
    let w1 = context.r1.wrapping_add(context.x[1]);
    let w2 = context.r2 ^ context.x[2];
    let tmp1 = l_1(concat_low_and_high(w1, w2));
    let tmp2 = l_2(concat_low_and_high(w2, w1));
    context.r1 = concat_u32(S0[(tmp1 >> 24) as usize], S1[((tmp1 >> 16) & 0xFF) as usize],
    S0[((tmp1 >> 8) & 0xFF) as usize], S1[(tmp1 & 0xFF) as usize]);
    context.r2 = concat_u32(S0[(tmp2 >> 24) as usize], S1[((tmp2 >> 16) & 0xFF) as usize],
    S0[((tmp2 >> 8) & 0xFF) as usize], S1[(tmp2 & 0xFF) as usize]);
    w
}

// · L || · H
fn concat_low_and_high(n1: u32, n2: u32) -> u32 {
    (n1 << 16) | (n2 >> 16)
}

fn l_1(x: u32) -> u32 {
    x ^ x.rotate_left(2) ^ x.rotate_left(10) ^ x.rotate_left(18) ^ x.rotate_left(24)
}

fn l_2(x: u32) -> u32 {
    x ^ x.rotate_left(8) ^ x.rotate_left(14) ^ x.rotate_left(22) ^ x.rotate_left(30)
}

fn lfsr_with_initialization_mode(u: u32, s: &mut[u32; 16]) {
    let next = lfsr_next(s);
    lfsr_append(s, add_mod(next, u));
}

fn lfsr_append(s: &mut[u32; 16], u: u32) {
    for i in 1..16 {
        s[i - 1] = s[i];
    }
    s[15] = u;
}

fn lfsr_next(s: &mut[u32; 16]) -> u32 {
    let mut f = s[0];
    let mut v = rotate_left_u31(s[0], 8);
    f = add_mod(f, v);
    v = rotate_left_u31(s[4], 20);
    f = add_mod(f, v);
    v = rotate_left_u31(s[10], 21);
    f = add_mod(f, v);
    v = rotate_left_u31(s[13], 17);
    f = add_mod(f, v);
    v = rotate_left_u31(s[15], 15);
    f = add_mod(f, v);
    f
}

fn rotate_left_u31(a: u32, shift: u32) -> u32 {
    ((a << shift) | ((a & 0x7fffffff) >> (31 - shift))) & 0x7FFFFFFF
}

// 模 2^31 - 1 加法
fn add_mod(a: u32, b: u32) -> u32 {
    let c = a.wrapping_add(b);
    (c & 0x7FFFFFFF) + (c >> 31)
}

fn lfsr_with_work_mode(s: &mut[u32; 16]) {
    let next = lfsr_next(s);
    lfsr_append(s, next);
}